cimport cython
import numpy
import scipy
cimport libc.stdlib

cimport numpy
numpy.import_array()
import scipy.optimize

DTYPE = numpy.float64
ctypedef numpy.float64_t DTYPE_t


cdef extern from "math.h":
    double log(double x)
    double exp(double x)

cpdef int pwLocate(DTYPE_t x, numpy.ndarray[DTYPE_t,ndim=1] s):
    cdef int N=len(s),j,found=-1
    for j in range(N-1):
        if x>=s[j] and x<=s[j+1]:
            found=j
            break
    return found

cdef DTYPE_t Heaviside(DTYPE_t x):
     if x>0: return 1
     else:   return 0

cdef DTYPE_t ErlangLikelyhood(numpy.ndarray[DTYPE_t,ndim=1] x,
                              numpy.ndarray[DTYPE_t,ndim=1] y,
                              numpy.ndarray[DTYPE_t,ndim=1] s,
                              numpy.ndarray[DTYPE_t,ndim=1] t0,
                              numpy.ndarray[DTYPE_t,ndim=1] a,
                              DTYPE_t R):
    cdef DTYPE_t mprod=0.0,dT
    cdef int N=len(x),M=len(s),i,j

    for i in range(N):
        j=pwLocate(x[i],s)
        if j==-1:
            mprod=log(1e-50)
            break
        else:
            #dT=y[i]-(x[i]/R+t0[j]+1.0/a[j])
            dT=y[i]-(x[i]/R+t0[j])
            if Heaviside(dT)>0:
                mprod+=log(Heaviside(dT)*dT*a[j]*a[j]*exp(-a[j]*dT))
            else:
                mprod+=log(1e-50)
    return mprod

cpdef DTYPE_t pErlangLikelyhood(numpy.ndarray[DTYPE_t,ndim=1] params,
                                numpy.ndarray[DTYPE_t,ndim=1] x,
                                numpy.ndarray[DTYPE_t,ndim=1] y,
                                numpy.ndarray[DTYPE_t,ndim=1] s):
    cdef int N=(len(params)-1)/2
    cdef numpy.ndarray[DTYPE_t,ndim=1] t0=params[1:N+1],a=params[N+1:2*N+1]
    return -ErlangLikelyhood(x,y,s,t0,a,params[0])

cpdef DTYPE_t GApErlangLikelyhood(numpy.ndarray[DTYPE_t,ndim=1] params,
                                numpy.ndarray[DTYPE_t,ndim=1] x,
                                numpy.ndarray[DTYPE_t,ndim=1] y,
                                numpy.ndarray[DTYPE_t,ndim=1] s):
    cdef int N=(len(params)-1)/2
    cdef numpy.ndarray[DTYPE_t,ndim=1] t0=params[1:N+1],a=params[N+1:2*N+1]
    return ErlangLikelyhood(x,y,s,t0,a,params[0])

cpdef numpy.ndarray[DTYPE_t,ndim=1] gExpectedFunc(numpy.ndarray[DTYPE_t,ndim=1] x,
                                          numpy.ndarray[DTYPE_t,ndim=1] s,
                                          numpy.ndarray[DTYPE_t,ndim=1] params):
    cdef int N=len(x),i,j,M=(len(params)-1)/2
    cdef numpy.ndarray[DTYPE_t,ndim=1] y=numpy.zeros(N)
    cdef numpy.ndarray[DTYPE_t,ndim=1] a=params[M+1:2*M+1],t0=params[1:M+1]
    cdef DTYPE_t  R=params[0]
    for i in range(N):
        j=pwLocate(x[i],s)
        if j==-1:y[i]=0.0
        else: y[i]=(x[i]/R+t0[j]+2.0/a[j])
    return y
    
cpdef numpy.ndarray[DTYPE_t,ndim=1] gMinFunc(numpy.ndarray[DTYPE_t,ndim=1] x,
                                          numpy.ndarray[DTYPE_t,ndim=1] s,
                                          numpy.ndarray[DTYPE_t,ndim=1] params):
    cdef int N=len(x),i,j,M=(len(params)-1)/2
    cdef numpy.ndarray[DTYPE_t,ndim=1] y=numpy.zeros(N)
    cdef numpy.ndarray[DTYPE_t,ndim=1] a=params[M+1:2*M+1],t0=params[1:M+1]
    cdef DTYPE_t  R=params[0]
    for i in range(N):
        j=pwLocate(x[i],s)
        if j==-1:y[i]=0.0
        else: y[i]=(x[i]/R+t0[j])
    return y
    
cpdef numpy.ndarray[DTYPE_t,ndim=1] gModeFunc(numpy.ndarray[DTYPE_t,ndim=1] x,
                                          numpy.ndarray[DTYPE_t,ndim=1] s,
                                          numpy.ndarray[DTYPE_t,ndim=1] params):
    cdef int N=len(x),i,j,M=(len(params)-1)/2
    cdef numpy.ndarray[DTYPE_t,ndim=1] y=numpy.zeros(N)
    cdef numpy.ndarray[DTYPE_t,ndim=1] a=params[M+1:2*M+1],t0=params[1:M+1]
    cdef DTYPE_t  R=params[0]
    for i in range(N):
        j=pwLocate(x[i],s)
        if j==-1:y[i]=0.0
        else: y[i]=(x[i]/R+t0[j]+1.0/a[j]) #t0[j]+ x[i]/R+1.0/a[j]
    return y
